(ns util.mon
  (:use clojure.tools.logging
        [util.exec :only [popen]]
        util.os
        util.utils
        ver)
  (:require clojure.string)
  (:import (org.hyperic.sigar Sigar NetFlags)
           (org.hyperic.sigar.ptql ProcessFinder)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; state
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defn mk-monbox []
  {:sigar (Sigar.)
   :cache (atom {})})

(def- thread-local-monbox (proxy [ThreadLocal] []
                            (initialValue [] (mk-monbox))))

(def ^:dynamic *monbox* nil)
;(def ^:dynamic *monbox* (mk-monbox))

(defn current-monbox []
  (or *monbox* (.get thread-local-monbox)))

(defn sigar []
  (:sigar (current-monbox)))

(defn- update-cached [k new-value]
  (let [a (:cache (current-monbox))
        old (atom nil)]
    (swap! a (fn [state] (reset! old (get state k)) (assoc state k new-value)))
    @old))

(defmacro in-monbox [monbox & body]
  `(binding [*monbox* ~monbox]
    ~@body))

(defmacro in-new-monbox [& body]
  `(in-monbox (mk-monbox)
    ~@body))



;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; ...
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defn get-core-number [] ; the Java way
  (.. Runtime getRuntime availableProcessors))

(defn get-core-number [] ; the SIGAR way
  (count (.getCpuList (sigar))))

(def get-core-number (memoize get-core-number))

(defn get-system-cpu []
  (.getCombined (.getCpuPerc (sigar))))

(defn get-wait-cpu []
  (let [v (.getWait (.getCpuPerc (sigar)))]
    (if (Double/isNaN v)
      0.0
      v)))

(defn get-system-memory []
  (.getTotal (.getMem (sigar))))

(defn get-system-memory-used []
  (.getActualUsed (.getMem (sigar))))

(defn get-system-memory-free []
  (.getActualFree (.getMem (sigar))))

(defn get-pid-list []
  (into [] (.getProcList (sigar))))

(defn get-process-cpu [pid]
  (/ (.getPercent (.getProcCpu (sigar) (long pid))) (get-core-number)))

(defn ^:dynamic get-process-memory [pid]
  " One Windows, the :memory is Working Set, and the :vm is Virtual Size"
  (let [mem (.getProcMem (sigar) (long pid))]
    {:vm (.getSize mem)
     :memory (.getResident mem)
     }))

(defn get-pid []
  (.getPid (sigar)))

(defn get-ppid [pid]
  (.getPpid (.getProcState (sigar) (long pid))))

(defn get-process-name [pid]
  (.getName (.getProcState (sigar) (long pid))))

(defn get-process-bin [pid]
  (.getName (.getProcExe (sigar) (long pid))))

(def get-process-exe get-process-bin)

(defn get-process-cwd [pid]
  (.getCwd (.getProcExe (sigar) (long pid))))

(defn get-process-args [pid]
  (into [] (.getProcArgs (sigar) (long pid))))

(defn get-process-env [pid]
  (into {} (.getProcEnv (sigar) (long pid))))

(defn get-process-info [pid]
  (let [pid (long pid)
        s (sigar)
        state (.getProcState s pid)
        exe (try (.getProcExe s pid) (catch Exception _ nil)) ; XXX: why fails in IDEA? because of the x86 JDK?
        env (try (.getProcEnv s pid) (catch Exception _ nil)) ; XXX: why fails in IDEA? because of the x86 JDK?
        mem (in-monbox (current-monbox) (get-process-memory pid))] ; use PDH API to read memory
    {:pid pid
     :ppid (.getPpid state)
     :name (.getName state)
     :exe (when exe (.getName exe)) ; XXX: why fails in IDEA? because of the x86 JDK?
     :bin (when exe (.getName exe)) ; XXX: why fails in IDEA?
     :args (try (into [] (.getProcArgs s pid)) (catch Exception _ nil)) ; XXX: why fails in IDEA?
     :cwd (when exe (.getCwd exe)) ; XXX: why fails in IDEA?
     :vm (:vm mem) ; use PDH API to read memory
     :memory (:memory mem) ; use PDH API to read memory
     :cpu (get-process-cpu pid)
     :env (when env (into {} env))}))

;; sometimes get-process-info throws "The parameter is incorrect" exception when trying to get-process-cpu
;; get-process-info-without-performance may avoid it...
(defn get-process-info-without-performance [pid]
  (let [pid (long pid)
        s (sigar)
        state (.getProcState s pid)
        exe (try (.getProcExe s pid) (catch Exception _ nil)) ; XXX: why fails in IDEA? because of the x86 JDK?
        env (try (.getProcEnv s pid) (catch Exception _ nil))] ; XXX: why fails in IDEA? because of the x86 JDK?
    {:pid pid
     :ppid (.getPpid state)
     :name (.getName state)
     :exe (when exe (.getName exe)) ; XXX: why fails in IDEA? because of the x86 JDK?
     :bin (when exe (.getName exe)) ; XXX: why fails in IDEA?
     :args (try (into [] (.getProcArgs s pid)) (catch Exception _ nil)) ; XXX: why fails in IDEA?
     :cwd (when exe (.getCwd exe)) ; XXX: why fails in IDEA?
     :env (when env (into {} env))}))

(defn sigar-find-pid-by-tcp-port [port] ; not supported on win 7. TODO: can we fix it using sigar source code?
  (let [pid (.getProcPort (sigar) NetFlags/CONN_TCP (long port))]
    (when (zero? pid)
      (throw (Exception. (str "No process is using TCP port " port))))
    pid))

(declare win-netstat-find-pid-by-tcp-port)
(defn find-pid-by-tcp-port [port]
  (try
    (sigar-find-pid-by-tcp-port port)
    (catch org.hyperic.sigar.SigarNotImplementedException e
      (if (windows?)
        (win-netstat-find-pid-by-tcp-port port)
        (throw e)))))

(defn find-pids-by-name [name]
  (into [] (.find (ProcessFinder. (sigar)) (str "State.Name.eq=" name))))

(defn find-pid-by-name [name]
  (let [finder (ProcessFinder. (sigar))
        q (str "State.Name.eq=" name)]
    (try
      (.findSingleProcess finder q)
      (catch org.hyperic.sigar.SigarException e
        (error e (format "Found %s for process %s" (seq (.find finder q)) name))
        (throw e)))))

(defn find-pid-by-pid-or-name [pid]
  (if (integer? pid)
    pid
    (try
      (Integer/parseInt pid)
      (catch NumberFormatException _
        (find-pid-by-name pid)))))

(defn process-exists? [pid]
  (if (integer? pid)
    (not (exception-thrown? (.getProcState (sigar) (long pid))))
    (seq (find-pids-by-name pid))))

(declare win-netstat-ports)
(defn used-ports []
  (if (windows?)
    (into {} (win-netstat-ports)) ; XXX: should it return a set, a map, of a list of port->pid pairs?
    (throw (Exception. "Not implemented"))))

(defn port-used? [port]
  (not (exception-thrown? (find-pid-by-tcp-port port))))

(defn ports-used? [ports]
  (not (every? #(exception-thrown? (find-pid-by-tcp-port %)) ports)))

(declare win-netstat-ports-used?)
(defn ports-used? [ports]
  (try
    (not (every? #(zero? (.getProcPort (sigar) NetFlags/CONN_TCP (long %))) ports))
    (catch org.hyperic.sigar.SigarNotImplementedException e
      (if (windows?)
        (win-netstat-ports-used? ports)
        (throw e)))))

(def SIGINT  2)
(def SIGKILL 9)
(def SIGTERM 15)
(defn kill-process
  ([pid signum]
    (let [pid (cond
                (integer? pid) (long pid)
                :else pid)
          signum (condp = signum
                   :default SIGTERM
                   :force   SIGKILL
                   :sigterm  SIGTERM
                   :sigkill SIGKILL
                   :sigint  SIGINT
                   signum)]
      (.kill (sigar) pid signum)))
  ([pid]
    (kill-process pid :default)))


;;; Windows specific operations

(defn win-netstat-ports []
  "Pairs of [port pid] -- note: not map"
  (let [output (popen ["netstat" "-ano"])
        lines (.split output "\r\n")
        lines (remove #(re-find #"\bFIN_WAIT_2\b" %) lines) ; ignore FIN_WAIT_2
        lines (map #(clojure.string/replace % #"^\s+|\s+$" "") (filter #(re-find #"^\s*TCP\b" %) lines))
        pp (for [line lines]
             (let [fields (clojure.string/split line #"\s+")
                   local (second fields)
                   port (Integer/parseInt (last (re-find #":(\d+)$" local)))
                   pid (Long/parseLong (last fields))]
               [port pid]))] ; XXX: should i remove 0 from pid list?
    pp))

(defn win-netstat-find-pid-by-tcp-port [port]
  (let [pp (win-netstat-ports)
        matched (set (map second (filter (fn [[p]] (= p port)) pp)))
        matched (disj matched 0 (long 0))] ; XXX: why remove 0?
    ;(debug output)
    (when (empty? matched)
      (throw (Exception. (str "No process is using TCP port " port))))
    (when (next matched)
      (throw (Exception. (format "More than one process is using TCP port %d: %s" port (seq matched))))) ; TODO: print process names
    (first matched)))

(defn win-netstat-used-ported []
  (let [pp (win-netstat-ports)
        used-ports (set (map first pp))]
    used-ports))

(defn win-netstat-ports-used? [ports]
  (some (set (win-netstat-used-ported)) ports))

(defn win-netstat-port-used? [port]
  (win-netstat-ports-used? [port]))

(defn pdh-get-process-memory [pid]
  (throw (Exception. (str "PDH is only available on Windows, but you are under " (get-os)))))
(defn with-pdh-box-fn [f]
  (f))
(defmacro with-pdh-box [& body]
  `(with-pdh-box-fn (fn [] ~@body)))
(defn winapi-get-process-memory [pid]
  (throw (Exception. (str "WINAPI is only available on Windows, but you are under " (get-os)))))
(defn pdh-update [])

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; disk io
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defn get-disk-list []
  (map #(.getDirName %) (filter #(= (.getType %) org.hyperic.sigar.FileSystem/TYPE_LOCAL_DISK) (.getFileSystemList (sigar)))))

(defn total-disk-io []
  (reduce +
    (for [disk (get-disk-list)]
      (let [fsu (.getFileSystemUsage (sigar) disk)] ; XXX: how to use getDiskUsage?
        (+ (.getDiskReadBytes fsu) (.getDiskWriteBytes fsu))))))

(defn get-disk-io []
  (let [current-time (System/currentTimeMillis)
        current-io (total-disk-io)
        [old-time old-io] (update-cached :disk [current-time current-io])]
    (if old-time
      (let [duration (/ (- current-time old-time) 1000)
            bytes (- current-io old-io)
            bytes (if (< bytes 0) (bit-and bytes 0xffffffff) bytes)]
        (long (/ bytes (if (zero? duration) 1 duration))))
      0)))

(defn get-disk-space [disk]
  (let [fsu (.getFileSystemUsage (sigar) disk)]
    {:total (* (.getTotal fsu) 1024)
     :free (* (.getAvail fsu) 1024)
     :used (* (.getUsed fsu) 1024)}))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; network io
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defn total-network-io []
  (reduce +
    (for [[addr items] (group-by first
                         (for [ni (.getNetInterfaceList (sigar))]
                           (let [stat (.getNetInterfaceStat (sigar) ni)
                                 addr (.getHwaddr (.getNetInterfaceConfig (sigar) ni))]
                             [addr (+ (.getRxBytes stat) (.getTxBytes stat))])))
          :when (not (= addr org.hyperic.sigar.NetFlags/NULL_HWADDR))]
      (long (/ (reduce + (map second items)) (count items))))))

(defn get-network-io []
  (let [current-time (System/currentTimeMillis)
        current-io (total-network-io)
        [old-time old-io] (update-cached :network [current-time current-io])]
    (if old-time
      (let [duration (/ (- current-time old-time) 1000)
            bytes (- current-io old-io)
            bytes (if (< bytes 0) (bit-and bytes 0xffffffff) bytes)]
        (long (/ bytes (if (zero? duration) 1 duration))))
      0)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; windows patch
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(when (windows?)
  (var-patch-vars 'util.mon 'util.mon.windows))

(comment
(println (.getNetStat (sigar)))
(println (find-pid-by-tcp-port 8080))
(println (.getTcp (sigar)))
(loop []
  (Thread/sleep 1000)
  (println (get-system-cpu))
  (recur))
)
